import traceback

from flask_sqlalchemy import SQLAlchemy
import time
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import OneHotEncoder

db = SQLAlchemy()


def get_time():
    return time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))


def dict_to_str(data):
    data = vars(data)
    for j in data.keys():
        data[j] = str(data[j])
    return data


def model_to_str_dict(data):
    return {str(k): str(v) for k, v in data.__dict__.items()}


def delete_page_key(item: dict):
    keys = [*item.keys()]
    if "page" in keys:
        del item["page"]
    # if "url" in keys:
    #     del item["url"]
    # if "_sa_instance_state" in keys:
    #     def item["_sa_instance_state"]
    return item


def delete_sa_instance_state(data):
    data = vars(data)
    del data['_sa_instance_state']
    return data


def get_html_by_selenium(url: str, timeout=60, cookies: dict = None) -> str:
    chrome_options = Options()
    # chrome_options.add_argument('--headless')
    driver = webdriver.Chrome(options=chrome_options)
    try:
        driver.get(url)
        if cookies:
            for k, v in cookies.items():
                driver.add_cookie({'name': k, 'value': v})
            driver.refresh()
        time.sleep(timeout)
        content = driver.page_source
    except Exception as e:
        traceback.print_exc()
        raise e
    finally:
        driver.close()
    return content


def predict_tea_price(tea_data, tea_name):
    # 读取数据并转换成DataFrame
    df = pd.DataFrame(tea_data, columns=['Name', 'Price'])

    # 使用独热编码对名称进行编码
    encoder = OneHotEncoder()
    name_encoded = encoder.fit_transform(df[['Name']])

    # 构建特征矩阵
    features = name_encoded

    # 构建目标向量
    target = df['Price']

    # 初始化线性回归模型
    model = LinearRegression()

    # 拟合模型
    model.fit(features, target)

    # 对新的名称进行编码
    name_encoded_new = encoder.transform([[tea_name]])

    # 使用模型进行预测
    predicted_price = model.predict(name_encoded_new)

    return predicted_price[0]
